#pragma once
#ifndef DWCONV_H_
#define DWCONV_H_

#include "stream_tools.h"

template<
	unsigned N_IO,		// # of IO
	unsigned N_CH,		// # of CHannel
	unsigned VEC_LEN	// ROW * COL
>
void dwconv_3x3(
	data_stream<N_IO * BIT_ACTV>& in, 
	data_stream<N_IO * BIT_CONV>& out, 
	const ap_uint<N_IO * BIT_WGHT> weight[N_CH / N_IO][9])
{
	static_assert(N_CH >= N_IO, "dwconv_3x3");
	static_assert(N_CH % N_IO == 0, "dwconv_3x3");
	constexpr unsigned FOLD = N_CH / N_IO;
	constexpr unsigned ITERS = VEC_LEN;

	assert(in.size() == VEC_LEN * FOLD * 9);
	assert(out.empty());

	#pragma HLS bind_storage variable=weight type=rom_1p impl=lutram
	ap_int<BIT_CONV> acc[N_IO];
	#pragma HLS array_partition variable=acc complete dim=1
	#pragma HLS bind_storage variable=acc type=ram_2p impl=lutram

	for (unsigned i = 0; i < N_IO; ++i)
	{
		#pragma HLS UNROLL
		acc[i] = 0;
	}

	for (unsigned it = 0; it < ITERS; ++it)
	{
		for (unsigned f = 0; f < FOLD; ++f)
		{
			for (unsigned k = 0; k < 9; ++k)
			{
				#pragma HLS PIPELINE II=1
				// load
				ap_uint<N_IO * BIT_ACTV> in_buf = in.read();
				ap_uint<N_IO * BIT_WGHT> wt_buf = weight[f][k];

				// calc
				for (unsigned i = 0; i < N_IO; ++i)
				{
					#pragma HLS UNROLL
					ap_uint<BIT_ACTV> x = in_buf(SLICE(BIT_ACTV, i));
					ap_int<BIT_WGHT>  y = wt_buf(SLICE(BIT_WGHT, i));
					acc[i] += x * y;
				}

				// output
				if (k == 8)
				{
					ap_uint<N_IO * BIT_CONV> out_buf;
					for (unsigned i = 0; i < N_IO; ++i)
					{
						#pragma HLS UNROLL
						out_buf(SLICE(BIT_CONV, i)) = acc[i];
						acc[i] = 0;
					}
					out.write(out_buf);
				}
			}
		}
	}

	assert(in.empty());
	assert(out.size() == VEC_LEN * FOLD);
	return;
}

template<
	unsigned N_IO,		// # of IO
	unsigned N_CH,		// # of CHannel
	unsigned N_K,		// # of Kernel
	unsigned N_G,		// # of Reduce
	unsigned VEC_LEN	// ROW * COL
>
void dwconv_3x3_r(
	data_stream<N_IO / N_G * BIT_ACTV>& in,
	data_stream<N_IO / N_G * BIT_CONV>& out,
	const ap_uint<N_IO / N_G * BIT_WGHT> weight[N_CH / N_IO][N_K * N_K][N_G])
{
	static_assert(N_CH >= N_IO, "dwconv_3x3");
	static_assert(N_CH % N_IO == 0, "dwconv_3x3");
	static_assert(N_IO % N_G == 0, "dwconv_3x3");
	constexpr unsigned FOLD = N_CH / N_IO;
	constexpr unsigned ITERS = VEC_LEN;

	assert(in.size() == VEC_LEN * FOLD * 9);
	assert(out.empty());

	#pragma HLS bind_storage variable=weight type=rom_1p impl=lutram
	ap_int<BIT_CONV*N_IO/N_G> acc[N_G];
	#pragma HLS bind_storage variable=acc type=ram_2p impl=lutram
	ap_int<BIT_CONV> y[N_IO/N_G];
	#pragma HLS array_partition variable=y complete dim=1

	for (unsigned g = 0; g < N_G; ++g)
	{
	#pragma HLS UNROLL
		acc[g] = 0;
	}

	for (unsigned it = 0; it < ITERS; ++it)
	{
		for (unsigned f = 0; f < FOLD; ++f)
		{
			for (unsigned k = 0; k < 9; ++k)
			{
				for (unsigned g = 0; g < N_G; ++g)
				{
					#pragma HLS PIPELINE II=1
					// load
					ap_uint<N_IO / N_G * BIT_ACTV> in_buf = in.read();
					ap_uint<N_IO / N_G * BIT_WGHT> wt_buf = weight[f][k][g];
					
					// calc
					for (unsigned i = 0; i < N_IO/N_G; ++i)
					{
						#pragma HLS UNROLL
						y[i] = acc[g](SLICE(BIT_CONV, i));
						ap_uint<BIT_ACTV> x = in_buf(SLICE(BIT_ACTV, i));
						ap_int<BIT_WGHT>  w = wt_buf(SLICE(BIT_WGHT, i));
						ap_int<BIT_CONV> wx = w * x;
						if(k == 0)
							y[i] = wx;
						else
							y[i] += wx;
						acc[g](SLICE(BIT_CONV, i)) = y[i];
					}

					// output
					if (k == 8)
					{
						ap_uint<N_IO / N_G * BIT_CONV> out_buf;
						for (unsigned i = 0; i < N_IO/N_G; ++i)
						{
							#pragma HLS UNROLL
							out_buf(SLICE(BIT_CONV, i)) = y[i];
						}
						out.write(out_buf);
					}
				}
			}
		}
	}

	assert(in.empty());
	assert(out.size() == VEC_LEN * FOLD);
	return;
}

#endif
